package poc

import (
	"os"
	"path/filepath"
	"strings"

	"gopkg.in/yaml.v3"
)

// POC 漏洞POC结构
type POC struct {
	ID          string            `yaml:"id"`
	Name        string            `yaml:"name"`
	Author      string            `yaml:"author"`
	Description string            `yaml:"description"`
	Severity    string            `yaml:"severity"` // critical, high, medium, low, info
	Tags        []string          `yaml:"tags"`
	Reference   []string          `yaml:"reference"`
	Request     Request           `yaml:"request"`
	Expression  string            `yaml:"expression"`
	Variables   map[string]string `yaml:"variables"`
}

// Request HTTP请求配置
type Request struct {
	Method  string            `yaml:"method"`
	Path    string            `yaml:"path"`
	Headers map[string]string `yaml:"headers"`
	Body    string            `yaml:"body"`
	FollowRedirects bool      `yaml:"follow_redirects"`
	MaxRedirects    int       `yaml:"max_redirects"`
}

// Response HTTP响应
type Response struct {
	StatusCode int
	Headers    map[string][]string
	Body       string
	Time       int64 // 响应时间(毫秒)
}

// LoadPOC 从文件加载POC
func LoadPOC(filename string) (*POC, error) {
	data, err := os.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	var poc POC
	if err := yaml.Unmarshal(data, &poc); err != nil {
		return nil, err
	}

	// 设置默认值
	if poc.Request.Method == "" {
		poc.Request.Method = "GET"
	}
	if poc.Request.MaxRedirects == 0 {
		poc.Request.MaxRedirects = 5
	}

	return &poc, nil
}

// LoadPOCsFromDir 从目录加载所有POC
func LoadPOCsFromDir(dir string) ([]*POC, error) {
	var pocs []*POC

	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// 只处理yaml和yml文件
		if !info.IsDir() && (strings.HasSuffix(path, ".yaml") || strings.HasSuffix(path, ".yml")) {
			poc, err := LoadPOC(path)
			if err != nil {
				// 记录错误但继续处理其他文件
				return nil
			}
			pocs = append(pocs, poc)
		}

		return nil
	})

	if err != nil {
		return nil, err
	}

	return pocs, nil
}

// GetSeverityLevel 获取严重程度数值
func (p *POC) GetSeverityLevel() int {
	switch strings.ToLower(p.Severity) {
	case "critical":
		return 5
	case "high":
		return 4
	case "medium":
		return 3
	case "low":
		return 2
	case "info":
		return 1
	default:
		return 0
	}
}